/***************************************
@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
@@!!!@@@@@@@@@@@@!!!!@@@@@@@@@@@%!!!@@@@@@@@@@@!!!!!!!!!!@@
@@!!!!!@@@@@@@@@%%!!!!!@@@@@@@@@%%!!!!@@@@@@@@@!!!!!!!!!!@@
@@@!!!!!@@@@@@@%%%%!!!!!@@@@@@@%%%!!!!!@@@@@@@@@@@@@!!!!!@@
@@@@!!!!!@@@@@@%%%@@!!!!!@@@@@%%%%@!!!!!@@@@@@@@@@@@!!!!!@@
@@@@@!!!!!@@@@%%%@@@@!!!!@@@@%%%%@@@!!!!!@@@@@@@@@@@!!!!!@@
@@@@@@!!!!!@@%%%%@@@@!!!!!@@@%%%@@@@@!!!!!@@@@@@@@@@!!!!!@@
@@@@@@!!!!!@%%%%@@@@@@!!!!!@%%%$@@@@@@!!!!@@@@@@@@@@!!!!!@@
@@@@@@@!!!!!%%%@@@@@@@@!!!!!%%%@@@@@@@@!!!!@@@@@@@@!!!!!!@@
@@@@@@@@!!!!!%@@@@@@@@@@!!!!!%@@@@@@@@@!!!!!!!!!!!!!!!!!!@@
@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@

WMJ战队深度学习目标检测 OpenVINO 推理后端
依赖库：                 OpenVINO-2022
目前支持使用的模型：       yolox改四点模型，大符/装甲板识别
更新时间：               2024.02.01
***************************************/

#pragma once
#include "InferenceBackend.hpp"
#include <openvino/openvino.hpp>

namespace wmj
{
    class OVBackend : public InferenceBackendBase
    {
    public: 
        OVBackend(){};
        OVBackend(const DetectorParams &params);
        ~OVBackend();
        bool doInference(const cv::Mat &inputImg);
        const float* getOutput();
        enum BackendType type_ = OV;

    private:
        ov::Tensor wrapMat2Tensor(const cv::Mat &mat);
        bool yolox_model_infer();

        ov::Core core;                        // 
        std::shared_ptr<ov::Model>  m_model;  // ov模型
        ov::CompiledModel m_compiled_model;   // 编译后模型
        ov::InferRequest m_infer_request;     // 推理请求

        float* m_prob;
        cv::Mat input_mat;                    // 图像输入
        ov::Tensor input_T;                   // 模型输入张量
        ov::Tensor output_T;                  // 模型输出张量
    };

    // Tools Functions
    ov::Tensor OVBackend::wrapMat2Tensor(const cv::Mat &mat)
    {
        const size_t channels = mat.channels();
        const size_t height = mat.size().height;
        const size_t width = mat.size().width;

        const size_t strideH = mat.step.buf[0];
        const size_t strideW = mat.step.buf[1];

        const bool is_dense = strideW == channels && strideH == channels * width;
        OPENVINO_ASSERT(is_dense, "Doesn't support conversion from not dense cv::Mat");

        return ov::Tensor(ov::element::u8, ov::Shape{1, height, width, channels}, mat.data);
    }

    OVBackend::OVBackend(const DetectorParams &params)
    {
        type_ = OV;
        assert(params.modelFormat == "onnx" || params.modelFormat == "xml");
        if(params.modelFormat == "onnx")
        {
            std::string model_name = MODEL_DIR + params.modelName + ".onnx";
            m_model = core.read_model(model_name);
        }
        else
        {
            std::string xml_file = MODEL_DIR + params.modelName + ".xml";
            std::string bin_file = MODEL_DIR + params.modelName + ".bin";
            m_model = core.read_model(xml_file,bin_file);
        }
        ov::preprocess::PrePostProcessor ppp(m_model);
        ov::preprocess::InputInfo &input_info = ppp.input();
        ov::preprocess::OutputInfo &output_info = ppp.output();
        input_info.tensor()
            .set_element_type(ov::element::u8)
            .set_layout("NHWC")
            .set_color_format(ov::preprocess::ColorFormat::BGR);
        input_info.model().set_layout("NCHW");
        if (params.modelType == "x" || params.modelType == "v8")
        {
            if(params.inputPrecision == "f32")
            {
                input_info.preprocess()
                            .convert_element_type(ov::element::f32);
            }
        }
        
        output_info.tensor().set_element_type(ov::element::f32);
        
        if(params.modelType == "v8")
        {
            output_info.model().set_layout("NHW");
            output_info.postprocess().convert_layout("NWH");
        }

        m_model = ppp.build();
        m_compiled_model = core.compile_model(m_model, params.DEVICE);
        m_infer_request = m_compiled_model.create_infer_request();
        input_T = m_infer_request.get_input_tensor();

        ov::Shape input_shape = m_model->input().get_shape();
        input_height = input_shape[1];
        input_width = input_shape[2];

        ov::Shape output_shape = m_model->output().get_shape();
        maxProposalCount = output_shape[1];
        objectLength = output_shape[2];
        input_mat = cv::Mat(input_height, input_width, CV_8UC3, cv::Scalar(105, 105, 105));
        return;
    }

    bool OVBackend::doInference(const cv::Mat &inputImg)
    {
        input_T = wrapMat2Tensor(inputImg);
        m_infer_request.set_input_tensor(input_T);

        m_infer_request.infer();

        output_T = m_infer_request.get_output_tensor();
        m_prob = output_T.data<float>();
        return true;
    }

    const float* OVBackend::getOutput()
    {
        return m_prob;
    }


}
